/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import Unsafe from './Unsafe';
import { StrategyInner, Strategy } from './CompressionParameters';
import BitOutputStream from './BitOutputStream';
import FseCompressionTable from './FseCompressionTable'
import SequenceStore from './SequenceStore'
import SequenceEncodingContext from './SequenceEncodingContext'
import NumberTransform from './NumberTransform'
import { FiniteStateEntropy } from './FiniteStateEntropy'
import Constants from './Constants'
import Histogram from './Histogram'
import Util from './Util'
import Exception from '../util/Exception'
import Long from '../util/long/index'

export default class SequenceEncoder {
    private static DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS_LOG: number = 6;
    private static DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS: Int16Array = new Int16Array([4, 3, 2, 2, 2, 2, 2, 2,
    2, 2, 2, 2, 2, 1, 1, 1,
    2, 2, 2, 2, 2, 2, 2, 2,
    2, 3, 2, 1, 1, 1, 1, 1,
    -1, -1, -1, -1]);
    private static DEFAULT_MATCH_LENGTH_NORMALIZED_COUNTS_LOG: number = 6;
    private static DEFAULT_MATCH_LENGTH_NORMALIZED_COUNTS: Int16Array = new Int16Array([1, 4, 3, 2, 2, 2, 2, 2,
    2, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, -1, -1,
    -1, -1, -1, -1, -1]);
    private static DEFAULT_OFFSET_NORMALIZED_COUNTS_LOG: number = 5;
    private static DEFAULT_OFFSET_NORMALIZED_COUNTS: Int16Array = new Int16Array([1, 1, 1, 1, 1, 1, 2, 2,
    2, 1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1, 1,
    -1, -1, -1, -1, -1]);
    private static DEFAULT_LITERAL_LENGTHS_TABLE: FseCompressionTable = FseCompressionTable.newInstance(SequenceEncoder.DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS, Constants.MAX_LITERALS_LENGTH_SYMBOL, SequenceEncoder.DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS_LOG);
    private static DEFAULT_MATCH_LENGTHS_TABLE: FseCompressionTable = FseCompressionTable.newInstance(SequenceEncoder.DEFAULT_MATCH_LENGTH_NORMALIZED_COUNTS, Constants.MAX_MATCH_LENGTH_SYMBOL, SequenceEncoder.DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS_LOG);
    private static DEFAULT_OFFSETS_TABLE: FseCompressionTable = FseCompressionTable.newInstance(SequenceEncoder.DEFAULT_OFFSET_NORMALIZED_COUNTS, Constants.DEFAULT_MAX_OFFSET_CODE_SYMBOL, SequenceEncoder.DEFAULT_OFFSET_NORMALIZED_COUNTS_LOG);

    constructor() {
    }

    public static compressSequences(outputBase: any, outputAddress: Long,
                                    outputSize: number, sequences: SequenceStore,
                                    strategy: StrategyInner, workspace: SequenceEncodingContext): number {
        let output: Long = outputAddress;
        let outputLimit: Long = outputAddress.add(outputSize);

        Util.checkArgument(outputLimit.sub(output).greaterThan(4) /* encoding type flags */
            , "Output buffer too small");

        let sequenceCount: number = sequences.sequenceCount;
        if (sequenceCount < 0x7F) {
            Unsafe.putByte(outputBase, output.toNumber(), NumberTransform.toByte(sequenceCount));
            output = output.add(1);
        }
        else if (sequenceCount < Constants.LONG_NUMBER_OF_SEQUENCES) {
            Unsafe.putByte(outputBase, output.toNumber(), NumberTransform.toByte(sequenceCount >>> 8 | 0x80));
            Unsafe.putByte(outputBase, output.add(1).toNumber(), NumberTransform.toByte(sequenceCount));
            output = output.add(Constants.SIZE_OF_SHORT);
        }
        else {
            Unsafe.putByte(outputBase, output.toNumber(), NumberTransform.toByte(0xFF));
            output = output.add(1);
            Unsafe.putShort(outputBase, output.toNumber(), NumberTransform.toShort(sequenceCount - Constants.LONG_NUMBER_OF_SEQUENCES));
            output = output.add(Constants.SIZE_OF_SHORT);
        }

        if (sequenceCount == 0) {
            return output.sub(outputAddress).toInt();
        }

        let headerAddress: Long = output;
        output = output.add(1);

        let maxSymbol: number;
        let largestCount: number;

        let counts: Int32Array = workspace.counts;
        Histogram.count(sequences.literalLengthCodes, sequenceCount, workspace.counts);
        maxSymbol = Histogram.findMaxSymbol(counts, Constants.MAX_LITERALS_LENGTH_SYMBOL);
        largestCount = Histogram.findLargestCount(counts, maxSymbol);

        let literalsLengthEncodingType: number = SequenceEncoder.selectEncodingType(largestCount, sequenceCount, SequenceEncoder.DEFAULT_LITERAL_LENGTH_NORMALIZED_COUNTS_LOG, true, strategy);

        let literalLengthTable: FseCompressionTable;
        switch (literalsLengthEncodingType) {
            case Constants.SEQUENCE_ENCODING_RLE:
                Unsafe.putByte(outputBase, output.toNumber(), sequences.literalLengthCodes[0]);
                output = output.add(1);
                workspace.literalLengthTable.initializeRleTable(maxSymbol);
                literalLengthTable = workspace.literalLengthTable;
                break;
            case Constants.SEQUENCE_ENCODING_BASIC:
                literalLengthTable = SequenceEncoder.DEFAULT_LITERAL_LENGTHS_TABLE;
                break;
            case Constants.SEQUENCE_ENCODING_COMPRESSED:
                output = output.add(SequenceEncoder.buildCompressionTable(
                    workspace.literalLengthTable,
                    outputBase,
                    output,
                    outputLimit,
                    sequenceCount,
                    Constants.LITERAL_LENGTH_TABLE_LOG,
                    sequences.literalLengthCodes,
                    workspace.counts,
                    maxSymbol,
                    workspace.normalizedCounts));
                literalLengthTable = workspace.literalLengthTable;
                break;
            default:
                throw new Exception("not yet implemented");
        }

        Histogram.count(sequences.offsetCodes, sequenceCount, workspace.counts);
        maxSymbol = Histogram.findMaxSymbol(counts, Constants.MAX_OFFSET_CODE_SYMBOL);
        largestCount = Histogram.findLargestCount(counts, maxSymbol);

        let defaultAllowed: boolean = maxSymbol < Constants.DEFAULT_MAX_OFFSET_CODE_SYMBOL;

        let offsetEncodingType: number = SequenceEncoder.selectEncodingType(largestCount, sequenceCount, SequenceEncoder.DEFAULT_OFFSET_NORMALIZED_COUNTS_LOG, defaultAllowed, strategy);

        let offsetCodeTable: FseCompressionTable;
        switch (offsetEncodingType) {
            case Constants.SEQUENCE_ENCODING_RLE:
                Unsafe.putByte(outputBase, output.toNumber(), sequences.offsetCodes[0]);
                output = output.add(1);
                workspace.offsetCodeTable.initializeRleTable(maxSymbol);
                offsetCodeTable = workspace.offsetCodeTable;
                break;
            case Constants.SEQUENCE_ENCODING_BASIC:
                offsetCodeTable = SequenceEncoder.DEFAULT_OFFSETS_TABLE;
                break;
            case Constants.SEQUENCE_ENCODING_COMPRESSED:
                output = output.add(SequenceEncoder.buildCompressionTable(
                    workspace.offsetCodeTable,
                    outputBase,
                    output,
                output.add(outputSize),
                    sequenceCount,
                    Constants.OFFSET_TABLE_LOG,
                    sequences.offsetCodes,
                    workspace.counts,
                    maxSymbol,
                    workspace.normalizedCounts));
                offsetCodeTable = workspace.offsetCodeTable;
                break;
            default:
                throw new Exception("not yet implemented");
        }

        Histogram.count(sequences.matchLengthCodes, sequenceCount, workspace.counts);
        maxSymbol = Histogram.findMaxSymbol(counts, Constants.MAX_MATCH_LENGTH_SYMBOL);
        largestCount = Histogram.findLargestCount(counts, maxSymbol);

        let matchLengthEncodingType: number = SequenceEncoder.selectEncodingType(largestCount, sequenceCount, SequenceEncoder.DEFAULT_MATCH_LENGTH_NORMALIZED_COUNTS_LOG, true, strategy);

        let matchLengthTable: FseCompressionTable;
        switch (matchLengthEncodingType) {
            case Constants.SEQUENCE_ENCODING_RLE:
                Unsafe.putByte(outputBase, output.toNumber(), sequences.matchLengthCodes[0]);
                output = output.add(1);
                workspace.matchLengthTable.initializeRleTable(maxSymbol);
                matchLengthTable = workspace.matchLengthTable;
                break;
            case Constants.SEQUENCE_ENCODING_BASIC:
                matchLengthTable = SequenceEncoder.DEFAULT_MATCH_LENGTHS_TABLE;
                break;
            case Constants.SEQUENCE_ENCODING_COMPRESSED:
                output = output.add(SequenceEncoder.buildCompressionTable(
                    workspace.matchLengthTable,
                    outputBase,
                    output,
                    outputLimit,
                    sequenceCount,
                    Constants.MATCH_LENGTH_TABLE_LOG,
                    sequences.matchLengthCodes,
                    workspace.counts,
                    maxSymbol,
                    workspace.normalizedCounts));
                matchLengthTable = workspace.matchLengthTable;
                break;
            default:
                throw new Exception("not yet implemented");
        }

        Unsafe.putByte(outputBase, headerAddress.toNumber(), (literalsLengthEncodingType << 6) | (offsetEncodingType << 4) | (matchLengthEncodingType << 2));

        output = output.add(SequenceEncoder.encodeSequences(outputBase, output, outputLimit, matchLengthTable, offsetCodeTable, literalLengthTable, sequences));

        return output.sub(outputAddress).toInt();
    }

    private static buildCompressionTable(table: FseCompressionTable, outputBase: object, output: Long, outputLimit: Long, sequenceCount: number, maxTableLog: number, codes: Int8Array, counts: Int32Array, maxSymbol: number, normalizedCounts: Int16Array): number {
        let tableLog: number = FiniteStateEntropy.optimalTableLog(maxTableLog, sequenceCount, maxSymbol);


        if (counts[codes[sequenceCount - 1]] > 1) {
            counts[codes[sequenceCount - 1]]--;
            sequenceCount--;
        }

        FiniteStateEntropy.normalizeCounts(normalizedCounts, tableLog, counts, sequenceCount, maxSymbol);
        table.initialize(normalizedCounts, maxSymbol, tableLog);

        return FiniteStateEntropy.writeNormalizedCounts(outputBase, output,
        outputLimit.sub(output).toInt(), normalizedCounts, maxSymbol, tableLog); // TODO: pass outputLimit directly
    }

    private static encodeSequences(
        outputBase: object,
        output: Long,
        outputLimit: Long,
        matchLengthTable: FseCompressionTable,
        offsetsTable: FseCompressionTable,
        literalLengthTable: FseCompressionTable,
        sequences: SequenceStore): number {
        let matchLengthCodes: Int8Array = sequences.matchLengthCodes;
        let offsetCodes: Int8Array = sequences.offsetCodes;
        let literalLengthCodes: Int8Array = sequences.literalLengthCodes;

        let blockStream: BitOutputStream = new BitOutputStream(outputBase, output, outputLimit.sub(output).toInt());

        let sequenceCount: number = sequences.sequenceCount;

        let matchLengthState: number = matchLengthTable.begin(matchLengthCodes[sequenceCount - 1]);
        let offsetState: number = offsetsTable.begin(offsetCodes[sequenceCount - 1]);
        let literalLengthState: number = literalLengthTable.begin(literalLengthCodes[sequenceCount - 1]);

        blockStream.addBits(sequences.literalLengths[sequenceCount - 1], Constants.LITERALS_LENGTH_BITS[literalLengthCodes[sequenceCount - 1]]);
        blockStream.addBits(sequences.matchLengths[sequenceCount - 1], Constants.MATCH_LENGTH_BITS[matchLengthCodes[sequenceCount - 1]]);
        blockStream.addBits(sequences.offsets[sequenceCount - 1], offsetCodes[sequenceCount - 1]);
        blockStream.flush();

        if (sequenceCount >= 2) {
            for (let n: number = sequenceCount - 2; n >= 0; n--) {
                let literalLengthCode: number = literalLengthCodes[n];
                let offsetCode: number = offsetCodes[n];
                let matchLengthCode: number = matchLengthCodes[n];

                let literalLengthBits: number = Constants.LITERALS_LENGTH_BITS[literalLengthCode];
                let offsetBits: number = offsetCode;
                let matchLengthBits: number = Constants.MATCH_LENGTH_BITS[matchLengthCode];

                offsetState = offsetsTable.encode(blockStream, offsetState, offsetCode); // 15
                matchLengthState = matchLengthTable.encode(blockStream, matchLengthState, matchLengthCode); // 24
                literalLengthState = literalLengthTable.encode(blockStream, literalLengthState, literalLengthCode); // 33

                if ((offsetBits + matchLengthBits + literalLengthBits >= 64 - 7 - (Constants.LITERAL_LENGTH_TABLE_LOG + Constants.MATCH_LENGTH_TABLE_LOG + Constants.OFFSET_TABLE_LOG))) {
                    blockStream.flush();
                }

                blockStream.addBits(sequences.literalLengths[n], literalLengthBits);
                if (((literalLengthBits + matchLengthBits) > 24)) {
                    blockStream.flush();
                }

                blockStream.addBits(sequences.matchLengths[n], matchLengthBits);
                if ((offsetBits + matchLengthBits + literalLengthBits > 56)) {
                    blockStream.flush();
                }

                blockStream.addBits(sequences.offsets[n], offsetBits); // 31
                blockStream.flush();
            }
        }

        matchLengthTable.finish(blockStream, matchLengthState);
        offsetsTable.finish(blockStream, offsetState);
        literalLengthTable.finish(blockStream, literalLengthState);

        let streamSize: number = blockStream.close();
        Util.checkArgument(streamSize > 0, "Output buffer too small");

        return streamSize;
    }

    private static selectEncodingType(
        largestCount: number,
        sequenceCount: number,
        defaultNormalizedCountsLog: number,
        isDefaultTableAllowed: boolean,
        strategy: StrategyInner): number {
        if (largestCount == sequenceCount) { // => all entries are equal
            if (isDefaultTableAllowed && sequenceCount <= 2) {
                return Constants.SEQUENCE_ENCODING_BASIC;
            }

            return Constants.SEQUENCE_ENCODING_RLE;
        }

        if (strategy.ordinal() < Strategy.LAZY.ordinal()) {
            if (isDefaultTableAllowed) {
                let factor = 10 - strategy.ordinal();
                let baseLog = 3;
                let minNumberOfSequences =
                Long.fromNumber(1)
                    .shiftLeft(defaultNormalizedCountsLog)
                    .multiply(factor)
                    .shiftRight(baseLog); /* 28-36 for offset, 56-72 for lengths */
                let belooms = (sequenceCount >> (defaultNormalizedCountsLog - 1))
                if (Long.fromNumber(sequenceCount).lessThan(minNumberOfSequences) || (largestCount < belooms)) {

                    return Constants.SEQUENCE_ENCODING_BASIC;
                }
            }
        } else {
            throw new Exception("not yet implemented");
        }
        return Constants.SEQUENCE_ENCODING_COMPRESSED;
    }
}
